import json
import ColorAndStyle
from drawTwoDim import DrawTwoDim


class DrawMDVRPResult:
    def __init__(self):
        self._nodes = []
        self._depot_count = 0
        self._customer_count = 0
        self._s = None

    def draw(self, problem_file, result_file, pic_file):
        self.read_instance(problem_file)
        self._s = DrawTwoDim(self._nodes)
        self.draw_nodes()
        self.read_lines(result_file)
        self._s.show(pic_file)

    def read_instance(self, filename):
        with open(filename, 'r') as f:
            line_num = 0
            for line in f:
                if line_num == 0:
                    numbers = line.split(' ')
                    self._customer_count = int(numbers[2])
                    self._depot_count = int(numbers[3])
                    self._nodes = [[0.0, 0.0] for _ in range(self._depot_count + self._customer_count)]
                elif line_num <= self._depot_count:
                    pass
                elif line_num < self._depot_count+self._customer_count+1:
                    numbers = ' '.join(line.split()).split(' ')
                    i = line_num-1
                    self._nodes[i][0] = float(numbers[1])
                    self._nodes[i][1] = float(numbers[2])
                else:
                    i = line_num-self._depot_count-self._customer_count-1
                    numbers = ' '.join(line.split()).split(' ')
                    self._nodes[i][0] = float(numbers[1])
                    self._nodes[i][1] = float(numbers[2])
                line_num += 1

    def draw_nodes(self):
        self._s.draw_nodes([x for x in range(self._depot_count)], [str(i) for i in range(self._depot_count)],
                           marksize=25)
        self._s.draw_nodes([x for x in range(self._depot_count, self._depot_count + self._customer_count)],
                           [str(i) for i in range(self._depot_count, self._depot_count + self._customer_count)],
                           color=ColorAndStyle.COLOR.YELLOW, marker=ColorAndStyle.STYLE.star)

    def read_lines(self, filename):
        with open(filename, 'r') as f:
            data = json.load(f)
            color_index = 0
            pre_depot = 0
            arcs = []
            for route in data['routes']:
                ids = route.strip().split(' ')
                for i in range(len(ids)):
                    ids[i] = int(ids[i])
                if pre_depot != ids[0]:
                    self._s.draw_lines(arcs, color=ColorAndStyle.COLOR.COLORS[color_index])
                    color_index += 1
                    if color_index == len(ColorAndStyle.COLOR.COLORS):
                        color_index = 0
                    arcs = []
                    pre_depot = ids[0]
                for i in range(len(ids)-1):
                    arcs.append([ids[i], ids[i+1]])
                arcs.append([ids[-1], ids[0]])
            self._s.draw_lines(arcs, color=ColorAndStyle.COLOR.COLORS[color_index])


if __name__ == "__main__":
    s = DrawMDVRPResult()
    s.draw("MDVRPInstance/pr10", 'p10.res', 'res')
